import jieba
from gensim.models import Word2Vec
from gensim.models.word2vec import PathLineSentences
from s01_train_vec import *
import config
from sklearn.decomposition import PCA

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt

# 设置字体为黑体，以支持中文
plt.rcParams['font.family'] = 'SimHei'  # 或者你安装的任何支持中文的字体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示为方块的问题
     

def run():

    model = load_model(config.model_file)

    # 假设model是你已经训练好的Word2Vec模型
    top_words = 300
    common_words = model.wv.index_to_key[:top_words]  # 获取前1000个最高频的词语
    common_vectors = model.wv[common_words]  # 获取这些词语的词向量



    # 初始化PCA模型，并将维度设置为3
    pca = PCA(n_components=3)

    # 对词向量进行降维
    common_vectors_reduced = pca.fit_transform(common_vectors)

    # 选择降维后的词向量的前三个维度
    x = common_vectors_reduced[:, 0]  # 第一维
    y = common_vectors_reduced[:, 1]  # 第二维
    z = common_vectors_reduced[:, 2]  # 第三维

    # 创建一个3D散点图
    fig = plt.figure(figsize=(16, 16))
    ax = fig.add_subplot(111, projection='3d')

    # 绘制点
    # ax.scatter(x, y, z)

    # 为每个点添加词语标签
    for i, word in enumerate(common_words):
        ax.text(x[i], y[i], z[i], word, size=10, zorder=1, color='k')

    ax.set_xlabel('X Label')
    ax.set_ylabel('Y Label')
    ax.set_zlabel('Z Label')

    # 设置x, y, z轴的范围
    ax.set_xlim(-2.5, 2.5)
    ax.set_ylim(-2.5, 2.5)
    ax.set_zlim(-2.5, 2.5)

    # 显示图形
    plt.show()


if __name__=='__main__':
    run()